from __future__ import annotations

from base64 import b64encode
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import TYPE_CHECKING, List

import binwalk
from pydantic import BaseModel, Field

import config
from analysis.plugin import AnalysisPluginV0
from analysis.plugin.plugin import AnalysisFailedError
from helperFunctions.install import OperateInDirectory
from plugins.mime_blacklists import MIME_BLACKLIST_COMPRESSED

if TYPE_CHECKING:
    import io

    from binwalk.modules.entropy import Entropy
    from binwalk.modules.signature import Signature


class SignatureScanResult(BaseModel):
    offset: int
    description: str


class AnalysisPlugin(AnalysisPluginV0):
    class Schema(BaseModel):
        entropy_analysis_graph: str = Field(
            description='An entropy analysis graph generated by binwalk as base64 string.',
        )
        signature_analysis: List[SignatureScanResult] = Field(
            description='The result of the signature analysis from binwalk.',
        )

    def __init__(self):
        super().__init__(
            metadata=AnalysisPluginV0.MetaData(
                name='binwalk',
                description='binwalk signature and entropy analysis',
                version='1.0.0',
                Schema=self.Schema,
                mime_blacklist=['audio/', 'image/', 'video/', 'text/', *MIME_BLACKLIST_COMPRESSED],
            ),
        )

    def analyze(self, file_handle: io.FileIO, virtual_file_path: dict[str, list[str]], analyses: dict) -> Schema:
        del virtual_file_path, analyses

        # FixMe: fix formatting once Python 3.8 is deprecated (2024-10-31)
        with TemporaryDirectory(
            prefix='fact_analysis_binwalk_', dir=config.backend.temp_dir_path
        ) as tmp_dir, OperateInDirectory(tmp_dir):
            output: tuple[Signature, Entropy] = binwalk.scan(
                file_handle.name,
                signature=True,
                entropy=True,
                save=True,
                quiet=True,
            )
            signature_result, entropy_result = output
            if not entropy_result.output_file or not (pic_path := Path(entropy_result.output_file)).is_file():
                raise AnalysisFailedError('Entropy output file is missing')
            return self.Schema(
                entropy_analysis_graph=b64encode(pic_path.read_bytes()).decode(),
                signature_analysis=[
                    SignatureScanResult(offset=i.offset, description=i.description) for i in signature_result.results
                ],
            )

    def summarize(self, result: Schema) -> list:
        summary = []
        for item in result.signature_analysis:  # type: SignatureScanResult
            if 'entropy edge' in item.description:
                continue
            if ',' in item.description:
                summary.append(item.description.split(',', maxsplit=1)[0])
            elif item.description:
                summary.append(item.description)
        return summary
